package com.ly.crawl.impl;

import com.ly.crawl.Capturer;
import org.jsoup.nodes.Document;
import org.jsoup.nodes.Element;
import org.jsoup.select.Elements;

import java.util.*;

public class SimpleUrlCapturer implements Capturer<Set<String>> {
    private String[] hrefRegexs;
    private String cssSelects;
    // 响应重定向映射
    private Map<String, String> redirects = new HashMap<>();

    public SimpleUrlCapturer(String cssSelects, String hrefRegex) {
        this.cssSelects = cssSelects == null ? "a[href]" : cssSelects;
        if (hrefRegex != null) {
            this.hrefRegexs = hrefRegex.split("[\\s,;，；]+");
            for (int i = 0; i < this.hrefRegexs.length; i++) {
                this.hrefRegexs[i] = this.hrefRegexs[i].replaceAll("\\d+", "\\\\d+");
            }
        }
    }

    public SimpleUrlCapturer redirectTo(String hrefPrefix, String redirectPrefix) {
        cssSelects += ",a[href^=" + hrefPrefix + "]";
        redirects.put(hrefPrefix, redirectPrefix);
        return this;
    }

    @Override
    public Set<String> capture(String url, Document doc) {
        Elements es = doc.select(cssSelects);
        if (es.isEmpty()) {
            return Collections.emptySet();
        } else {
            Set<String> ret = new HashSet<>();
            Iterator<Element> iterator = es.iterator();
            while (iterator.hasNext()) {
                String href = iterator.next().attr("href");
                if (!matchHrefRegexs(href)) {
                    continue;
                }
                href = castHref(href);
                ret.add(href);
            }
            return ret;
        }
    }

    private boolean matchHrefRegexs(String href) {
        if (this.hrefRegexs == null) {
            return true;
        }
        for (String hrefRegex : this.hrefRegexs) {
            if (href.matches(hrefRegex)) {
                return true;
            }
        }
        return false;
    }

    private String castHref(final String href) {
        String hrefPrefix = redirects.keySet().stream().filter(h -> href.startsWith(h)).findFirst().orElse(null);
        if (hrefPrefix == null) {
            return href;
        } else {
            String redirectHref = redirects.get(hrefPrefix);
            return href.replace(hrefPrefix, redirectHref);
        }
    }
}
